{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Execute code in a sandbox\n",
    "\n",
    "To enhance security and protect yourself from malicious code through prompt injection, \n",
    "we make it possible to run code in a sandbox environment.\n",
    "This notebook explains how to do it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install the package\n",
    "\n",
    "First of all you need to install the python package. You can use pip to install it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install pandasai-docker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Execute the code in the sandbox\n",
    "\n",
    "Please keep in mind the sandbox works offline. \n",
    "Once you have installed the package, you can start the sandbox with the following code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandasai as pai\n",
    "from pandasai_docker import DockerSandbox\n",
    "from pandasai_litellm.litellm import LiteLLM\n",
    "\n",
    "# Initialize LiteLLM with your OpenAI model\n",
    "llm = LiteLLM(model=\"gpt-4.1-mini\", api_key=\"YOUR_OPENAI_API_KEY\")\n",
    "\n",
    "# Configure PandasAI to use this LLM\n",
    "pai.config.set({\n",
    "    \"llm\": llm\n",
    "})\n",
    "\n",
    "# initialize the sandbox\n",
    "sandbox = DockerSandbox()\n",
    "sandbox.start()\n",
    "\n",
    "# read a csv as df\n",
    "df = pai.read_csv(\"./data/heart.csv\")\n",
    "\n",
    "# pass the csv and the sandbox to the agent\n",
    "result = pai.chat(\"plot total heart patients by gender\", df, sandbox=sandbox)\n",
    "\n",
    "result.show()\n",
    "\n",
    "# stop the sandbox (docker container)\n",
    "sandbox.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Execute the code in the sandbox with the agent\n",
    "\n",
    "Please keep in mind the sandbox works offline. \n",
    "Once you have installed the package, you can start the sandbox with the following code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandasai as pai\n",
    "from pandasai import Agent\n",
    "from pandasai_docker import DockerSandbox\n",
    "from pandasai_litellm.litellm import LiteLLM\n",
    "\n",
    "# Initialize LiteLLM with your OpenAI model\n",
    "llm = LiteLLM(model=\"gpt-4.1-mini\", api_key=\"YOUR_OPENAI_API_KEY\")\n",
    "\n",
    "# Configure PandasAI to use this LLM\n",
    "pai.config.set({\n",
    "    \"llm\": llm\n",
    "})\n",
    "\n",
    "# initialize the sandbox\n",
    "sandbox = DockerSandbox()\n",
    "sandbox.start()\n",
    "\n",
    "# read a csv as df\n",
    "df = pai.read_csv(\"./data/heart.csv\")\n",
    "\n",
    "# pass the csv and the sandbox to the agent\n",
    "agent = Agent([df], memory_size=10, sandbox=sandbox)\n",
    "\n",
    "# Chat with the Agent\n",
    "response = agent.chat(\"plot top five artists streams\")\n",
    "\n",
    "# stop the sandbox (docker container)\n",
    "sandbox.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Customize the sandbox\n",
    "\n",
    "You can decide the name and path of your sandbox by passing them as positional arguments in the DockerSandbox()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sandbox = DockerSandbox(\"PandasAI-sandbox\", \"/path/to/Dockerfile\")\n",
    "\n",
    "# read a csv as df\n",
    "df = pai.read_csv(\"./data/heart.csv\")\n",
    "\n",
    "# pass the csv and the sandbox to the agent\n",
    "agent = Agent([df], memory_size=10, sandbox=sandbox)\n",
    "\n",
    "# Chat with the Agent\n",
    "response = agent.chat(\"plot top five artists streams\")\n",
    "\n",
    "sandbox.stop()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
